package sampler;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.apache.http.HttpResponse;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.concurrent.FutureCallback;
import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
import org.apache.http.impl.nio.client.HttpAsyncClients;
import org.apache.http.impl.nio.conn.PoolingNHttpClientConnectionManager;
import org.apache.http.impl.nio.reactor.DefaultConnectingIOReactor;
import org.apache.http.nio.reactor.ConnectingIOReactor;
import org.apache.http.util.EntityUtils;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;
import org.apache.jmeter.samplers.SampleResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class MyAsyncSampler implements JavaSamplerClient {

	private final Logger log = LoggerFactory.getLogger(MyAsyncSampler.class);

	public int MaxTotal;
	public int DefaultMaxPerRoute;
	public int SocketTimeout;
	public int ConnectTimeout;
	public int ConnectionRequestTimeout;

	public String MaxTotal_TEMPLATE_VALUE = "150";
	public String DefaultMaxPerRoute_TEMPLATE_VALUE = "150";
	public String SocketTimeout_TEMPLATE_VALUE = "250";
	public String ConnectTimeout_TEMPLATE_VALUE = "250";
	public String ConnectionRequestTimeout_TEMPLATE_VALUE = "1000";

	public ConnectingIOReactor ioReactor;
	public PoolingNHttpClientConnectionManager connectionManager;
	public CloseableHttpAsyncClient httpclient;
	public RequestConfig requestConfig;


	public List<HttpGet> requests = new ArrayList<HttpGet>();
	public String REQUESTS_TEMPLATE_VALUE = "[]";


	@Override
	public Arguments getDefaultParameters() {

		Arguments params = new Arguments();

		params.addArgument("SetMaxTotal", MaxTotal_TEMPLATE_VALUE, "metadata", "");
		params.addArgument("SetDefaultMaxPerRoute", DefaultMaxPerRoute_TEMPLATE_VALUE, "metadata", "");
		params.addArgument("SetSocketTimeout", SocketTimeout_TEMPLATE_VALUE, "metadata", "");
		params.addArgument("SetConnectTimeout", ConnectTimeout_TEMPLATE_VALUE, "metadata", "");
		params.addArgument("SetConnectionRequestTimeout", ConnectionRequestTimeout_TEMPLATE_VALUE, "metadata", "");

		params.addArgument("SetRequests", REQUESTS_TEMPLATE_VALUE, "metadata", "");

		return params;
	}

	@SuppressWarnings("unused")
	@Override
	public void setupTest(JavaSamplerContext context) {

		MaxTotal = context.getIntParameter("SetMaxTotal", 150);
		DefaultMaxPerRoute = context.getIntParameter("SetDefaultMaxPerRoute", 150);
		SocketTimeout = context.getIntParameter("SetSocketTimeout", 250);
		ConnectTimeout = context.getIntParameter("SetConnectTimeout", 250);
		ConnectionRequestTimeout = context.getIntParameter("SetConnectionRequestTimeout", 1000);

		log.info("AsyncSampler [MaxTotal=" + MaxTotal + ", DefaultMaxPerRoute=" + DefaultMaxPerRoute
				+ ", SocketTimeout=" + SocketTimeout + ", ConnectTimeout=" + ConnectTimeout
				+ ", ConnectionRequestTimeout=" + ConnectionRequestTimeout + "]");

		try {

			ioReactor = new DefaultConnectingIOReactor();
			connectionManager = new PoolingNHttpClientConnectionManager(ioReactor);
			connectionManager.setMaxTotal(MaxTotal);
			connectionManager.setDefaultMaxPerRoute(DefaultMaxPerRoute);

			httpclient = HttpAsyncClients.custom().setConnectionManager(connectionManager).build();
			httpclient.start();

			requestConfig = RequestConfig.custom().setSocketTimeout(SocketTimeout).setConnectTimeout(ConnectTimeout)
					.setConnectionRequestTimeout(ConnectionRequestTimeout).build();

		} catch (Exception e) {
			// TODO: handle exception
			e.printStackTrace();
		}

		JSONArray tempArray = JSON.parseArray(context.getParameter("SetRequests", "[]"));
		for (Object object : tempArray) {

			JSONObject tempObject = (JSONObject) object;

			String method = tempObject.getString("method");
			String url = tempObject.getString("url");
			String headers = tempObject.getString("headers");
			String body = tempObject.getString("body");

			switch (method.toLowerCase()) {
				case "get":
					HttpGet tempGet = new HttpGet(url);
					tempGet.setConfig(requestConfig);
//				tempGet.addHeader("", "");
					requests.add(tempGet);
					break;

				default:
					break;
			}
		}
	}

	@Override
	public SampleResult runTest(JavaSamplerContext context) {


		CountDownLatch countDownLatch = new CountDownLatch(requests.size());


		StringBuffer buffer = new StringBuffer();

		SampleResult sampleResult = new SampleResult();

		sampleResult.sampleStart();

		sampleResult.setSuccessful(true);

		for (HttpGet get : requests) {

			httpclient.execute(get, new FutureCallback<HttpResponse>() {

				@Override
				public void cancelled() {
					// TODO Auto-generated method stub
					countDownLatch.countDown();
				}

				@Override
				public void completed(HttpResponse response) {
					try {
						buffer.append(EntityUtils.toString(response.getEntity(), "UTF-8"));
						buffer.append("\n\n");
					} catch (Exception e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					} finally {
						countDownLatch.countDown();
					}
				}

				@Override
				public void failed(Exception exception) {
					// TODO Auto-generated method stub
					sampleResult.setSuccessful(false);
					countDownLatch.countDown();
				}

			});
		}

		try {
			countDownLatch.await();
		} catch (InterruptedException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

		sampleResult.setResponseData(buffer.toString(), "utf-8");

		sampleResult.sampleEnd();

		return sampleResult;
	}

	@Override
	public void teardownTest(JavaSamplerContext context) {

	}
}
